/*
hipcc singleProcess.cpp -o singleProcess -lrccl
./singleProcess
*/
#include <cstdio>
#include <hip/hip_runtime.h>
#include <iostream>
#include <rccl/rccl.h>
#include <string>
#define HIP_CALL(cmd)                                                                                                  \
    do {                                                                                                               \
        hipError_t error = (cmd);                                                                                      \
        if (error != hipSuccess) {                                                                                     \
            std::cerr << "Encountered HIP error (" << hipGetErrorString(error) << ") at line" << __LINE__              \
                      << " in file " << __FILE__ << "\n";                                                              \
            exit(-1);                                                                                                  \
        }                                                                                                              \
    } while (0)
#define NCCL_CALL(cmd)                                                                                                 \
    do {                                                                                                               \
        ncclResult_t error = (cmd);                                                                                    \
        if (error != ncclSuccess) {                                                                                    \
            std::cerr << "Encountered NCCL error (" << ncclGetErrorString(error) << ") at line " << __LINE__           \
                      << " in file " << __FILE__ << "\n";                                                              \
            exit(-1);                                                                                                  \
        }                                                                                                              \
    } while (0)
int main() {
    int numIntraRank = 4;
    ncclComm_t comm[numIntraRank];
    // 初始化通信组内所有的rank
    NCCL_CALL(ncclCommInitAll(comm, numIntraRank, NULL));
    // Test configuration settings
    int minPow = 10;       // Starting power of 2 input size
    int maxPow = 28;       // Ending power of 2 input size
    int numIterations = 10;// Number of timed iterations
    // Allocate GPU resources for this process
    hipStream_t stream[numIntraRank];
    for (int i = 0; i < numIntraRank; i++) {
        HIP_CALL(hipSetDevice(i));
        HIP_CALL(hipStreamCreate(&stream[i]));
    }
    // Loop over power-of-two input sizes
    for (int power = minPow; power <= maxPow; power++) {
        int N = 1 << power;
        // Allocate GPU memory
        float *iputGpu[numIntraRank], *oputGpu[numIntraRank];
        for (int r = 0; r < numIntraRank; r++) {
            HIP_CALL(hipSetDevice(r));
            HIP_CALL(hipMalloc((void **) &iputGpu[r], N * sizeof(float)));
            HIP_CALL(hipMalloc((void **) &oputGpu[r], N * sizeof(float)));
        }
        // Allocate CPU memory for input/output
        float *iputCpu = (float *) malloc(N * sizeof(float));
        float *oputCpu = (float *) malloc(N * sizeof(float));
        // Fill CPU memory with a simple pattern
        for (int i = 0; i < N; i++) {
            iputCpu[i] = 1.0f;
            oputCpu[i] = 0.0f;
        }
        // Copy the input from CPU memory to GPU memory
        for (int r = 0; r < numIntraRank; r++) {
            HIP_CALL(hipSetDevice(r));
            HIP_CALL(hipMemcpy(iputGpu[r], iputCpu, N * sizeof(float), hipMemcpyHostToDevice));
        }
        NCCL_CALL(ncclGroupStart());
        for (int r = 0; r < numIntraRank; r++) {
            HIP_CALL(hipSetDevice(r));
            NCCL_CALL(ncclAllReduce(iputGpu[r], oputGpu[r], N, ncclFloat, ncclSum, comm[r], stream[r]));
        }
        NCCL_CALL(ncclGroupEnd());
        for (int r = 0; r < numIntraRank; r++) HIP_CALL(hipStreamSynchronize(stream[r]));
        printf("finished, count: %d\n", power);
        // Release GPU resources
        for (int r = 0; r < numIntraRank; r++) {
            HIP_CALL(hipFree(oputGpu[r]));
            HIP_CALL(hipFree(iputGpu[r]));
        }
        free(iputCpu);
        free(oputCpu);
    }
    for (int r = 0; r < numIntraRank; r++) {
        HIP_CALL(hipStreamDestroy(stream[r]));
        NCCL_CALL(ncclCommDestroy(comm[r]));
    }
}